#include "WebSocketServer.h"
#include <iostream>
#include <cstring>

static std::map<lws *, std::string> connections;

WebSocketServer::WebSocketServer(int port) : port(port), context(nullptr), running(false) {}

WebSocketServer::~WebSocketServer() {
    stop();
}

void WebSocketServer::start() {
    if (running) {
        std::cerr << "Server is already running" << std::endl;
        return;
    }

static struct lws_protocols protocols[] = {
    {
        "http",                           // 协议名称
        callback_websockets,              // 回调函数
        0,                                // 每个会话的数据大小
        1024,                             // 接收缓冲区大小
        0,                                // 协议 ID
        NULL                              // 用户数据指针
    },
    { NULL, NULL, 0, 0, 0, NULL }       // 终结符
};

    struct lws_context_creation_info info;
    memset(&info, 0, sizeof(info));

    info.port = port;
    info.protocols = protocols;
    info.gid = -1;
    info.uid = -1;

    context = lws_create_context(&info);
    if (context == NULL) {
        std::cerr << "Failed to create libwebsocket context" << std::endl;
        return;
    }

    running = true;
    server_thread = std::thread(&WebSocketServer::run, this);
}

void WebSocketServer::stop() {
    if (!running) {
        return;
    }

    {
        std::lock_guard<std::mutex> lock(context_mutex);
        lws_context_destroy(context);
        context = nullptr;
    }

    running = false;
    if (server_thread.joinable()) {
        server_thread.join();
    }
}

void WebSocketServer::run() {
    while (running) {
        lws_service(context, 50);
    }
}


int WebSocketServer::callback_websockets(struct lws *wsi,
                                         enum lws_callback_reasons reason,
                                         void *user, void *in, size_t len) {

    switch (reason) {
        case LWS_CALLBACK_SERVER_WRITEABLE: {
            if (connections.find(wsi) != connections.end()) {
                std::string message = connections[wsi];
                unsigned char buf[LWS_SEND_BUFFER_PRE_PADDING + message.size() + LWS_SEND_BUFFER_POST_PADDING];
                unsigned char *p = &buf[LWS_SEND_BUFFER_PRE_PADDING];

                memcpy(p, message.data(), message.size());

                lws_write(wsi, p, message.size(), LWS_WRITE_TEXT);
                connections.erase(wsi);
            }
            break;
        }
        case LWS_CALLBACK_RECEIVE: {
            std::string message(reinterpret_cast<char*>(in), len);
            std::cout << "Received message: " << message << std::endl;
            
            connections[wsi] = message;

            lws_callback_on_writable( wsi);
            break;
        }
        case LWS_CALLBACK_ESTABLISHED: {
            std::cout << "New connection established" << std::endl;
            break;
        }
        case LWS_CALLBACK_CLOSED: {
            std::cout << "Connection closed" << std::endl;
            connections.erase(wsi);
            break;
        }
        default:
            break;
    }
    return 0;
}

